import os
import pandas as pd
import mysql.connector
import json
from datetime import datetime, timedelta
from process_chat_history import get_chat_time_bounds

def connect_to_database():
    """连接到数据库"""
    # 数据库配置
    db_config = {
        'host': os.getenv('MYSQL_HOST', 'localhost'),
        'user': os.getenv('MYSQL_USER', 'root'),
        'password': os.getenv('MYSQL_PASSWORD', 'neo..6555894'),
        'database': os.getenv('MYSQL_DATABASE', 'sales_chat_mvp')
    }
    
    try:
        conn = mysql.connector.connect(**db_config)
        print("数据库连接成功")
        return conn
    except Exception as e:
        print(f"数据库连接失败: {e}")
        return None

def get_all_records(conn):
    """获取所有需要更新的记录"""
    cursor = conn.cursor(dictionary=True)
    
    try:
        # 获取所有记录
        cursor.execute("""
        SELECT id, nickname, chat_history, week_start_time, time_span
        FROM chat_history
        """)
        
        records = cursor.fetchall()
        print(f"获取到 {len(records)} 条记录")
        return records
    except Exception as e:
        print(f"获取记录失败: {e}")
        return []
    finally:
        cursor.close()

def update_record_times(conn, record):
    """更新单条记录的时间信息"""
    cursor = conn.cursor()
    
    try:
        # 解析聊天历史
        chat_history = json.loads(record['chat_history'])
        
        if not chat_history:
            print(f"记录 {record['id']} 没有有效的聊天历史，跳过更新")
            return False
        
        # 创建临时DataFrame来分析时间
        # 首先构建包含发送人和消息内容的数据
        data = []
        for msg in chat_history:
            sender = list(msg.keys())[0]
            content = msg[sender]
            data.append({'发送人': sender, '消息内容': content})
        
        # 假设消息是按时间顺序排列的，我们用索引模拟时间
        # 由于我们只关心相对顺序，所以可以用当前时间加上索引作为模拟的时间
        base_time = datetime.now()
        message_times = [base_time + timedelta(minutes=i) for i in range(len(data))]
        
        # 创建DataFrame
        temp_df = pd.DataFrame(data)
        temp_df['消息时间'] = message_times
        
        # 获取实际的聊天开始和结束时间（这里是模拟的时间）
        start_time, end_time = get_chat_time_bounds(temp_df)
        
        # 获取现有的week_start_time作为基准时间
        base_week_start = record['week_start_time']
        
        # 由于我们使用的是模拟时间，我们只需要将相对差值应用到原始的week_start_time上
        # 这样我们就可以保持相对的开始和结束时间
        
        # 解析现有的time_span
        time_span = json.loads(record['time_span'])
        old_span_start = datetime.strptime(time_span['start_time'], '%Y-%m-%d %H:%M:%S')
        old_span_end = datetime.strptime(time_span['end_time'], '%Y-%m-%d %H:%M:%S')
        
        # 计算新的time_span
        # start_time应该与week_start_time相同
        # end_time保持原来的结束日期，但时间应该与chat_end_time一致
        new_time_span = {
            'start_time': base_week_start.strftime('%Y-%m-%d %H:%M:%S'),
            'end_time': old_span_end.strftime('%Y-%m-%d %H:%M:%S')
        }
        
        # 更新记录
        cursor.execute("""
        UPDATE chat_history 
        SET end_time = %s, time_span = %s
        WHERE id = %s
        """, (
            end_time.strftime('%Y-%m-%d %H:%M:%S'),
            json.dumps(new_time_span, ensure_ascii=False),
            record['id']
        ))
        
        print(f"更新记录 {record['id']} 成功")
        return True
    except Exception as e:
        print(f"更新记录 {record['id']} 失败: {e}")
        return False
    finally:
        cursor.close()

def process_actual_chat_history(conn):
    """处理实际聊天记录文件来更新数据库"""
    cursor = conn.cursor(dictionary=True)
    
    try:
        # 获取所有记录
        cursor.execute("""
        SELECT id, nickname, chat_history, week_start_time, time_span
        FROM chat_history
        """)
        
        records = cursor.fetchall()
        print(f"获取到 {len(records)} 条记录")
        
        # 聊天记录目录
        chat_dir = 'historychat'
        
        # 获取所有Excel文件
        excel_files = [f for f in os.listdir(chat_dir) if f.endswith('.xlsx') and not f.startswith('.~')]
        
        file_data = {}
        
        # 读取所有文件
        for file in excel_files:
            file_path = os.path.join(chat_dir, file)
            filename = os.path.basename(file_path)
            
            try:
                # 尝试使用不同的引擎读取
                try:
                    df = pd.read_excel(file_path, engine='openpyxl')
                except Exception:
                    df = pd.read_excel(file_path, engine='xlrd')
                
                # 确保时间列是datetime类型
                df['消息时间'] = pd.to_datetime(df['消息时间'], errors='coerce')
                df = df.dropna(subset=['消息时间'])
                
                if not df.empty:
                    # 从文件名提取昵称
                    parts = filename.split(' - ')
                    if len(parts) >= 2:
                        nickname = ' - '.join(parts[1:])
                        nickname = nickname.replace('.xlsx', '')
                        
                        # 存储数据
                        file_data[nickname] = df
                        print(f"已读取文件 {filename}, 昵称: {nickname}")
            except Exception as e:
                print(f"读取文件 {filename} 失败: {e}")
        
        # 更新记录
        updated_count = 0
        for record in records:
            nickname = record['nickname']
            
            if nickname in file_data:
                df = file_data[nickname]
                
                # 获取实际的聊天开始和结束时间
                start_time, end_time = get_chat_time_bounds(df)
                
                # 解析现有的time_span
                time_span = json.loads(record['time_span'])
                old_span_end = datetime.strptime(time_span['end_time'], '%Y-%m-%d %H:%M:%S')
                
                # 更新time_span
                new_time_span = {
                    'start_time': start_time.strftime('%Y-%m-%d %H:%M:%S'),
                    'end_time': old_span_end.strftime('%Y-%m-%d %H:%M:%S')
                }
                
                update_cursor = conn.cursor()
                
                try:
                    # 将datetime对象转换为字符串，避免MySQL类型转换错误
                    start_time_str = start_time.strftime('%Y-%m-%d %H:%M:%S')
                    end_time_str = end_time.strftime('%Y-%m-%d %H:%M:%S')
                    
                    # 更新记录
                    update_cursor.execute("""
                    UPDATE chat_history 
                    SET week_start_time = %s, end_time = %s, time_span = %s
                    WHERE id = %s
                    """, (
                        start_time_str,
                        end_time_str,
                        json.dumps(new_time_span, ensure_ascii=False),
                        record['id']
                    ))
                    
                    updated_count += 1
                    print(f"更新记录 {record['id']} 成功")
                except Exception as e:
                    print(f"更新记录 {record['id']} 失败: {e}")
                finally:
                    update_cursor.close()
            else:
                print(f"未找到昵称 '{nickname}' 对应的文件数据，跳过记录 {record['id']}")
        
        print(f"总计更新 {updated_count} 条记录")
    except Exception as e:
        print(f"处理聊天记录失败: {e}")
    finally:
        cursor.close()

def main():
    """主函数"""
    print("开始更新数据库中的时间字段...")
    
    # 连接数据库
    conn = connect_to_database()
    if not conn:
        print("无法连接到数据库，退出程序")
        return
    
    try:
        # 使用实际的聊天历史文件更新记录
        process_actual_chat_history(conn)
        
        # 提交事务
        conn.commit()
        print("所有更新已完成并提交到数据库")
    
    except Exception as e:
        conn.rollback()
        print(f"更新过程中出错，已回滚: {e}")
    
    finally:
        conn.close()
        print("数据库连接已关闭")

if __name__ == "__main__":
    main() 